#!/usr/bin/python
# -*- coding: utf-8 -*-
# Description: PE 工具，可以提取PE的头信息，版本信息，数字签名等
# Create: 2008-6-25
# Author: MK2[fengmk2@gmail.com]
import os
import time
from datetime import datetime
import struct
import codecs 
import pefile
from xml.dom.minidom import Document
from seblog.common_logic.utility.const import FileType, FileInfoKey, FixedFileInfoKey
#from fileauth.public.define.md5mgr import mksha1fromfile, mkmd5fromfile
from seblog.common_logic.utility import util 

#PE的常用扩展名
PE_EXT = set(('.exe', '.dll', '.sys', '.com', '.msi', '.dos'))

IMAGE_DOS_SIGNATURE = pefile.IMAGE_DOS_SIGNATURE
IMAGE_OS2_SIGNATURE = pefile.IMAGE_OS2_SIGNATURE
IMAGE_OS2_SIGNATURE_LE = pefile.IMAGE_OS2_SIGNATURE_LE
IMAGE_VXD_SIGNATURE = pefile.IMAGE_VXD_SIGNATURE
IMAGE_NT_SIGNATURE = pefile.IMAGE_NT_SIGNATURE

def is_pe_ext(filename):
    """检测是否以pe扩展名结束"""
    ext = os.path.splitext(filename)[1].lower()
    if ext and ext in PE_EXT:
        return True
    return False

def is_pe(filepath):
    """判断一个文件是否PE
    return (True/False, filetype)
    filetype: also see 'fileauth.public.define.const.FileType'
    
    >>> path = 'D:/testsamples/'
    >>> (True, FileType.EXE) == is_pe(path + 'exe/uplive.exe')
    True
    >>> (False, FileType.Empty) == is_pe(path + 'empty/blank.txt')
    True
    >>> (False, FileType.Unknow) == is_pe(path + 'unknow/2byte.txt')
    True
    """
    if not os.path.exists(filepath) or not os.path.isfile(filepath):
        return False, FileType.Nofile
    if os.path.getsize(filepath) == 0:
        return False, FileType.Empty
    type = FileType.Unknow
    try:
        f = open(filepath, u'rb')
        try:
            try:
                dos_signature = struct.unpack('H', f.read(2))[0]
                if dos_signature != IMAGE_DOS_SIGNATURE:
                    return False, type
                f.seek(60, 0)
                offset = struct.unpack('L', f.read(4))[0]
                f.seek(offset, 0)
                signature = struct.unpack('L', f.read(4))[0]
                result = False
                print signature
                if signature == IMAGE_NT_SIGNATURE:
                    result = True
                else:
                    f.seek(offset, 0)
                    signature = struct.unpack('H', f.read(2))[0]
                    if signature == IMAGE_OS2_SIGNATURE or signature == IMAGE_OS2_SIGNATURE_LE:
                        result = True
                if result:
                    f.seek(0, 0)
                    f.seek(offset+22, 0)
                    characteristics = struct.unpack('H', f.read(2))[0]
                    if characteristics >> 13 & 1 == 1:
                        type = FileType.DLL
                    else:
                        type = FileType.EXE
                return result, type
            except struct.error:
                return False, FileType.Unknow
        finally:
            f.close()
    except IOError, e:
        if e.errno == 13: #Permission denied
            return False, FileType.Permdenied
        raise 

VERSIONINFO_KEYS = set((
              FileInfoKey.Comments,
              FileInfoKey.CompanyName,
              FileInfoKey.FileDescription,
              FileInfoKey.FileVersion,
              FileInfoKey.InternalName,
              FileInfoKey.LegalCopyright,
              FileInfoKey.LegalTrademarks,
              FileInfoKey.ProductName, 
              FileInfoKey.ProductVersion, 
              FileInfoKey.PrivateBuild, 
              FileInfoKey.OriginalFilename,
              FileInfoKey.SpecialBuild,
              FileInfoKey.AppliesTo,
              FileInfoKey.BuildDate,
              FileInfoKey.InstallationType,
              FileInfoKey.InstallerEngine,
              FileInfoKey.InstallerVersion,
              FileInfoKey.KBArticleNumber,
              FileInfoKey.PackageType,
              FileInfoKey.ProcArchitecture,
              FileInfoKey.SelfExtractorVersion,
              FileInfoKey.SupportLink,
              FileInfoKey.Language,
              FileInfoKey.LanguageName,
              FileInfoKey.CodePage,
              FileInfoKey.CodePageName,))

VS_FIXEDFILEINFO_KEYS = set((
              FixedFileInfoKey.Signature,
              FixedFileInfoKey.StrucVersion,
              FixedFileInfoKey.FileVersionMS,
              FixedFileInfoKey.FileVersionLS,
              FixedFileInfoKey.ProductVersionMS,
              FixedFileInfoKey.ProductVersionLS,
              FixedFileInfoKey.FileFlagsMask,
              FixedFileInfoKey.FileFlags,
              FixedFileInfoKey.FileOS,
#              FixedFileInfoKey.FileType,
              FixedFileInfoKey.FileSubtype,
              FixedFileInfoKey.FileDateMS,
              FixedFileInfoKey.FileDateLS,))

def __get_versioninfo_from_peobject(pe):
    """get version info from pefile.PE object"""
    versioninfo = {FileInfoKey.HasVersionInfo: False}
    if hasattr(pe.FILE_HEADER, u'IMAGE_FILE_32BIT_MACHINE'):
        versioninfo[FileInfoKey.Attribute] = 32
    elif hasattr(pe.FILE_HEADER, u'IMAGE_FILE_16BIT_MACHINE'):
        versioninfo[FileInfoKey.Attribute] = 16
    try:
        createdtime = datetime.utcfromtimestamp(pe.FILE_HEADER.TimeDateStamp).strftime(u'%Y-%m-%d %X')
        if createdtime.year > 1980: # 避免日期是错误的....
            versioninfo[FileInfoKey.CreatedTime] = createdtime
    except:
        pass
    if hasattr(pe, u'SignedData'):
        versioninfo[FileInfoKey.DigitalIssuer] = pe.SignedData.Certificates[-1].TBSCertificate.Issuer
        versioninfo[FileInfoKey.DigitalSubject] = pe.SignedData.Certificates[-1].TBSCertificate.Subject
        versioninfo[FileInfoKey.DigitalStartDate] = \
            pe.SignedData.Certificates[-1].TBSCertificate.Validity.NotBefore.strftime('%Y-%m-%d %X')
        versioninfo[FileInfoKey.DigitalEndDate] = \
            pe.SignedData.Certificates[-1].TBSCertificate.Validity.NotAfter.strftime('%Y-%m-%d %X')
    if hasattr(pe, u'VS_FIXEDFILEINFO'):
        for keys in pe.VS_FIXEDFILEINFO.keys:
                for key in keys:
                    if key in VS_FIXEDFILEINFO_KEYS:
                        value = getattr(pe.VS_FIXEDFILEINFO, key)
                        if value is not None:
                            versioninfo[key] = value
    if pe.versioninfo:
        versioninfo[FileInfoKey.HasVersionInfo] = True
        for key, value in pe.versioninfo.iteritems():
            if key in VERSIONINFO_KEYS:
                versioninfo[key] = value
    return versioninfo

def get_pe_fileinfo(filepath, pe=None):
    """获取VersionInfo信息, 具体key可见VERSIONINFO_KEYS
    return dict contain version info.
    数字签名有关项:{
            DigitalSubject:       颁发者
            DigitalIssuer:        主题
            DigitalStartDate:     开始时间
            DigitalEndDate:       结束时间 }
    
    >>> path = 'D:/'
    >>> KPEanalyse_exe = get_pe_fileinfo(path + 'KPEanalyse.exe')
    >>> KPEanalyse_exe['FileDescription'], KPEanalyse_exe['SHA1']
    (u'KPEanaly \u5e94\u7528\u7a0b\u5e8f', '9d8c79e2edf156d2d731a1878b7b362a75ed0e3e')
    >>> KPEanalyse_exe['FileSize']
    21504L
    >>> KPEanalyse_exe['FileVersionMS']
    65538L
    >>> KPEanalyse_exe['ProductName']
    u' KPEanaly \u5e94\u7528\u7a0b\u5e8f'
    >>> KPEanalyse_exe['HasVersionInfo']
    True
    >>> KPEanalyse_exe['Language'] == 0x0804
    True
    >>> KPEanalyse_exe['CodePage']
    1200
    >>> KPEanalyseDll_dll = get_pe_fileinfo(path + 'KPEanalyseDll.dll')
    >>> KPEanalyseDll_dll['HasVersionInfo']
    True
    >>> KPEanalyseDll_dll['Language']
    2052
    >>> KPEanalyseDll_dll['Attribute']
    32
    >>> KPEanalyseDll_dll['CreatedTime']
    '2008-06-25 07:57:12'
    >>> get_pe_fileinfo(path + '/testsamples/pe2xml/qq.exe')['DigitalSubject'].split(os.linesep)
    ['C = US', 'O = VeriSign, Inc.', 'CN = VeriSign Time Stamping Services Signer - G2']
    """
    filestat = os.stat(filepath)
    fileinfo = {FileInfoKey.FileName: os.path.basename(filepath),
                FileInfoKey.FileSize: filestat.st_size,
                FileInfoKey.CreatedTime: datetime.utcfromtimestamp(filestat.st_ctime).strftime('%Y-%m-%d %X'),
                FileInfoKey.ModifiedTime: datetime.utcfromtimestamp(filestat.st_mtime).strftime('%Y-%m-%d %X'),
                FileInfoKey.AccessedTime: datetime.utcfromtimestamp(filestat.st_atime).strftime('%Y-%m-%d %X'),
                }
    fileinfo[FileInfoKey.HasVersionInfo], fileinfo[FileInfoKey.FileType] = is_pe(filepath)
    if fileinfo[FileInfoKey.HasVersionInfo]:
        if pe is None:
            pe = pefile.PE(filepath)
            try:
                pe.parse()
            except Exception, e:
                pe.close()
                raise 
        fileinfo.update(__get_versioninfo_from_peobject(pe))
        pe.close()
    return fileinfo

def __ensure_unicode(val):
    if isinstance(val, str):
        val = util.ensure_unicode(val)
    elif not isinstance(val, unicode):
        val = unicode(val)
    return val#util.ensure_unicode(val)
    
def __add_textnode(text, tagname, doc, parentnode):
    text = __ensure_unicode(text)
    tagname = tagname.replace(u' ', u'')
    node = doc.createElement(tagname)
    textnode = doc.createTextNode(text)
    node.appendChild(textnode)
    parentnode.appendChild(node)

def __set_attr(key, val, node):
    val = __ensure_unicode(val)
    key = key.replace(u' ', u'')
    node.setAttribute(key, val)
    
def __add_node_attrs_from_structure(structure, tagname, doc, parentnode):
    if tagname is None:
        tagname = structure.name
    node = doc.createElement(tagname)
    for keys in structure.keys:
        for key in keys:
            val = getattr(structure, key)
            if isinstance(val, int) or isinstance(val, long):
                if key == u'TimeDateStamp' or key == u'dwTimeStamp':
                    try:
                        val = time.strftime(u'%Y-%m-%d %X', time.gmtime(val))
                    except pefile.exceptions.ValueError, e:
                        val = u'%s [INVALID TIME]' % val
            __set_attr(key, val, node)
    parentnode.appendChild(node)
    return node

def __add_node_from_structure(structure, tagname, doc, parentnode):
    if tagname is None:
        tagname = structure.name
    node = doc.createElement(tagname)
    for keys in structure.keys:
        for key in keys:
            val = getattr(structure, key)
            if isinstance(val, int) or isinstance(val, long):
                if key == u'TimeDateStamp' or key == u'dwTimeStamp':
                    try:
                        val = time.strftime(u'%Y-%m-%d %X', time.gmtime(val))
                    except pefile.exceptions.ValueError, e:
                        val = u'%s [INVALID TIME]' % val
            __add_textnode(val, key, doc, node)
    parentnode.appendChild(node)
    return node

def __add_node_attrs_from_dict(d, tagname, doc, parentnode):
    node = doc.createElement(tagname)
    keys = d.keys()
    keys.sort()
    for key in keys:
        __set_attr(key, d[key], node)
    parentnode.appendChild(node)
    return node

def __add_node_from_dict(d, tagname, doc, parentnode):
    node = doc.createElement(tagname)
    keys = d.keys()
    keys.sort()
    for key in keys:
        __add_textnode(d[key], key, doc, node)
    parentnode.appendChild(node)
    return node

def __add_attrs_from_dict(dict, node):
    for key, value in dict.iteritems():
        __set_attr(key, value, node)

def __get_flags(container, flaglist):
    flags = {}
    for flag in flaglist:
        if getattr(container, flag[0]):
            flags[flag[0]] = 1 #只显示有的
    return flags

def pe2xml(filepath, xmlpath=None, pe=None):
    """保存PE信息到xml文件, 若不指定xmlpath，则默认为filepath + .xml为保存路径。
    return pefile.PE or None
    如果指定的文件是pe并且头信息是正确的，则返回pefile.PE，并生成xml；
    否则返回None，不生成xml
    """
#    assert is_pe(filepath)[0] == True, u'%s 不是PE.' % filepath
    if pe is None: #bad pe
        pe = pefile.PE(filepath)
        try:
            pe.parse()
        except pefile.PEFormatError, e:
            print e
            pe.close()
            return None
        except Exception, e:
            print e
            pe.close()
            raise 
#    warnings = pe.get_warnings()
#    if warnings: #bad pe
#        warnings_text = ''.join(['%s%s' % (w, os.linesep) for w in warnings])
#        return None
    if xmlpath is None:
        xmlpath = filepath + u'.xml'
    doc = Document()
    # Create the <File /> base element
    root = doc.createElement(u'File')
    doc.appendChild(root)
    # FileInfo
    try:
        fileinfo = get_pe_fileinfo(filepath, pe)

    except pefile.PEFormatError, e:
        print e
        pe.close()
        return None
#    if u'MD5' not in fileinfo: # 怎么算大文件的md5这么耗时间？
#        fileinfo[u'MD5'] = mkmd5fromfile(filepath)
#    if u'SHA1' not in fileinfo:
#        fileinfo[u'SHA1'] = mksha1fromfile(filepath)
    __add_node_from_dict(fileinfo, u'FileInfo', doc, root)
    
    # Warnings
#    warnings = doc.createElement(u'Warnings')
#    root.appendChild(warnings)
    if pe.warnings:
        waring_text = u''
        for warning in pe.warnings:
            waring_text += warning + os.linesep
        __add_textnode(waring_text, u'Warnings', doc, root)
    
    # PE
    PE = doc.createElement(u'PE')
    root.appendChild(PE)
    
    # PE/DOS_HEADER
    __add_node_from_structure(pe.DOS_HEADER, u'DOSHeader', doc, PE)
    # PE/NT_HEADERS
    __add_node_from_structure(pe.NT_HEADERS, u'NTHeader', doc, PE)
    # PE/FILE_HEADER
    file_header = __add_node_from_structure(pe.FILE_HEADER, u'FileHeader', doc, PE)
    image_flags = pe.retrieve_flags(pefile.IMAGE_CHARACTERISTICS, u'IMAGE_FILE_')
    
    flags = __get_flags(pe.FILE_HEADER, image_flags)
    # PE/OPTIONAL_HEADER/Flags
    __add_node_from_dict(flags, u'Flags', doc, file_header)
    # PE/OPTIONAL_HEADER
    if hasattr(pe, u'OPTIONAL_HEADER') and pe.OPTIONAL_HEADER is not None:
        optional_header = __add_node_from_structure(pe.OPTIONAL_HEADER, u'OptionalHeader', doc, PE)

        dll_characteristics_flags = pe.retrieve_flags(pefile.DLL_CHARACTERISTICS, u'IMAGE_DLL_CHARACTERISTICS_')
        flags = __get_flags(pe.OPTIONAL_HEADER, dll_characteristics_flags)
         # PE/OPTIONAL_HEADER/DllCharacteristics
        __add_node_from_dict(flags, u'DllCharacteristics', doc, optional_header)
    
    # PE/Sections
    sections = doc.createElement(u'Sections')
    PE.appendChild(sections)
    section_flags = pe.retrieve_flags(pefile.SECTION_CHARACTERISTICS, u'IMAGE_SCN_')
    for section in pe.sections:
        # PE/Sections/IMAGE_SECTION_HEADER
        section_node = __add_node_from_structure(section, u'SectionHeader', doc, sections)
        flags = __get_flags(section, section_flags)
        # PE/Sections/IMAGE_SECTION_HEADER/Flags
        __add_node_from_dict(flags, u'Flags', doc, section_node)
#        __add_textnode('%f (Min=0.0, Max=8.0)' % section.get_entropy(),
#                       'Entropy', doc, section_node)
#        if pefile.md5 is not None:
#            __add_textnode('%s' % section.get_hash_md5(),
#                       'MD5', doc, section_node)
#        if pefile.sha1 is not None:
#            __add_textnode('%s' % section.get_hash_sha1(),
#                       'SHA1', doc, section_node)
#        if pefile.sha256 is not None:
#            __add_textnode('%s' % section.get_hash_sha256(),
#                       'SHA256', doc, section_node)
#        if pefile.sha512 is not None:
#            __add_textnode('%s' % section.get_hash_sha512(),
#                       'SHA512', doc, section_node)
    # PE/IMAGE_DATA_DIRECTORY
    if (hasattr(pe, u'OPTIONAL_HEADER') and 
        hasattr(pe.OPTIONAL_HEADER, u'DATA_DIRECTORY')):
        directories = doc.createElement(u'DataDirectory')
        PE.appendChild(directories)
        for idx in xrange(len(pe.OPTIONAL_HEADER.DATA_DIRECTORY)):
            directory = pe.OPTIONAL_HEADER.DATA_DIRECTORY[idx]
            if directory.Size > 0:
                __add_node_attrs_from_structure(directory, None, doc, directories)
    
    # PE/VS_VERSIONINFO
    if hasattr(pe, u'VS_VERSIONINFO'):
        # PE/VS_VERSIONINFO
        versioninfo_node = __add_node_attrs_from_structure(pe.VS_VERSIONINFO, None, doc, PE)
        # PE/VS_VERSIONINFO/VS_FIXEDFILEINFO
        if hasattr(pe, u'VS_FIXEDFILEINFO'):
            __add_node_from_structure(pe.VS_FIXEDFILEINFO, None, 
                                      doc, versioninfo_node)
        
        # PE/VS_VERSIONINFO/StringFileInfo or VarFileInfo
        if hasattr(pe, u'FileInfo'):
            for entry in pe.FileInfo:
                fileinfo_node = __add_node_attrs_from_structure(entry, None, doc, versioninfo_node)
                
                # PE/VS_VERSIONINFO/StringFileInfo/StringTable
                if hasattr(entry, u'StringTable'):
                    for st_entry in entry.StringTable:
                        strtable_node = __add_node_attrs_from_structure(st_entry, None, doc, fileinfo_node)
                        __set_attr(u'LangID', st_entry.LangID, strtable_node)
                        for str_entry in st_entry.entries.items():
                            __add_textnode(str_entry[1], str_entry[0], doc, strtable_node)
                # PE/VS_VERSIONINFO/VarFileInfo/Var
                elif hasattr(entry, u'Var'):
                    for var_entry in entry.Var:
                        if hasattr(var_entry, u'entry'):
                            var_node = __add_node_attrs_from_structure(var_entry, None, doc, fileinfo_node)
                            __add_textnode(var_entry.entry.values()[0], var_entry.entry.keys()[0], 
                                           doc, var_node)
    # PE/IMAGE_DIRECTORY_ENTRY_EXPORT
    if hasattr(pe, u'DIRECTORY_ENTRY_EXPORT'):
        export_table_node = __add_node_attrs_from_structure(pe.DIRECTORY_ENTRY_EXPORT.struct, 
                                             u'ExportTable', doc, PE)
        for export in pe.DIRECTORY_ENTRY_EXPORT.symbols:
            # PE/IMAGE_DIRECTORY_ENTRY_EXPORT/Function
            if export.address is None:
                export.address = 0
            function = {u'Ordinal': u'%-10d' % export.ordinal, 
                        u'RVA': u'0x%08Xh' % export.address, 
                        u'Name': export.name
                      }
            if export.forwarder:
                function[u'Forwarder'] = export.forwarder
            __add_node_attrs_from_dict(function, u'Function', doc, export_table_node)
        
    if hasattr(pe, u'DIRECTORY_ENTRY_IMPORT'):
        # PE/IMAGE_DIRECTORY_ENTRY_IMPORT
        import_table_attrs = {u'NumberOfDLLs': len(pe.DIRECTORY_ENTRY_IMPORT)}
        import_table_node = __add_node_attrs_from_dict(import_table_attrs, 
                                                       u'ImportTable', doc, PE)
        functions_count = 0
        for module in pe.DIRECTORY_ENTRY_IMPORT:
            functions_count += len(module.imports)
            # PE/IMAGE_DIRECTORY_ENTRY_IMPORT/IMAGE_IMPORT_DESCRIPTOR
            module_node = __add_node_attrs_from_structure(module.struct, 
                                      None, 
                                      doc, import_table_node)
            __add_attrs_from_dict({u'DLL': module.dll,
                                   u'NumberOfFunctions': len(module.imports)
                                   }, module_node)
            for function in module.imports:
                # PE/IMAGE_DIRECTORY_ENTRY_IMPORT/IMAGE_IMPORT_DESCRIPTOR/Function
                function_attrs = {u'Name': function.name,
                                  u'Address': function.address,
                                  u'Hint': function.hint}
                if function.import_by_ordinal is True:
                    function_attrs[u'Name'] = \
                        u'Ordinal[%s] (Imported by Ordinal)' % function.ordinal
                else:
                    function_attrs[u'Name'] = function.name
                    function_attrs[u'Hint'] = function.hint
                if function.bound:
                    function_attrs[u'Bound'] = function.bound
                __add_node_attrs_from_dict(function_attrs, u'Function', doc, module_node)
        __set_attr(u'NumberOfFunctions', functions_count, import_table_node)
            
    # PE/IMAGE_DIRECTORY_ENTRY_BOUND_IMPORT
    if hasattr(pe, u'DIRECTORY_ENTRY_BOUND_IMPORT'):
        bound_imports_node = doc.createElement(u'BoundImportTable')
        PE.appendChild(bound_imports_node)
        for bound_imp_desc in pe.DIRECTORY_ENTRY_BOUND_IMPORT:
            
            # PE/DIRECTORY_ENTRY_BOUND_IMPORT/IMAGE_BOUND_IMPORT_DESCRIPTOR
            bound_imp_desc_node = __add_node_attrs_from_structure(bound_imp_desc.struct, 
                                      None, 
                                      doc, bound_imports_node)
            __set_attr(u'DLL', bound_imp_desc.name, bound_imp_desc_node)
            
            for bound_imp_ref in bound_imp_desc.entries:
                bound_imp_ref_node = __add_node_attrs_from_structure(bound_imp_ref.struct, 
                                      None, 
                                      doc, bound_imp_desc_node)
                __set_attr(u'DLL', bound_imp_ref.name, bound_imp_ref_node)
                
    # PE/IMAGE_DIRECTORY_ENTRY_DELAY_IMPORT
    if hasattr(pe, u'DIRECTORY_ENTRY_DELAY_IMPORT'):
        delay_imports_node = doc.createElement(u'DelayImportTable')
        PE.appendChild(delay_imports_node)
        functions_count = 0
        for module in pe.DIRECTORY_ENTRY_DELAY_IMPORT:
            functions_count += len(module.imports)
            # PE/IMAGE_DIRECTORY_ENTRY_DELAY_IMPORT/IMAGE_DELAY_IMPORT_DESCRIPTOR
            module_node = __add_node_attrs_from_structure(module.struct, 
                                      module.struct.name, 
                                      doc, delay_imports_node)
            __add_attrs_from_dict({u'DLL': module.dll, u'NumberOfFunctions': len(module.imports)}, 
                                  module_node)
            for function in module.imports:
                # PE/IMAGE_DIRECTORY_ENTRY_DELAY_IMPORT/IMAGE_DELAY_IMPORT_DESCRIPTOR/Function
                function_attrs = {u'Name': function.name,
                                  u'Address': function.address,
                                  u'Hint': function.hint}
                if function.import_by_ordinal is True:
                    function_attrs[u'Name'] = \
                        u'Ordinal[%s] (Imported by Ordinal)' % function.ordinal
                else:
                    function_attrs[u'Name'] = function.name
                    function_attrs[u'Hint'] = function.hint
                if function.bound:
                    function_attrs[u'Bound'] = function.bound
                __add_node_attrs_from_dict(function_attrs, u'Function', doc, module_node)
        __set_attr(u'NumberOfFunctions', functions_count, delay_imports_node)
                
    if hasattr(pe, u'DIRECTORY_ENTRY_RESOURCE'):
        # PE/IMAGE_DIRECTORY_ENTRY_RESOURCE
        resource_directory_node = __add_node_attrs_from_structure(
                                        pe.DIRECTORY_ENTRY_RESOURCE.struct, 
                                        u'Resource', doc, PE)
        for resource_type in pe.DIRECTORY_ENTRY_RESOURCE.entries:
            # PE/IMAGE_DIRECTORY_ENTRY_RESOURCE/Entry
            if resource_type.name is not None:
                resource_type_name = resource_type.name
            else:
                #resource_type_attrs['Type'] = u'0x%X (%s)' % \ #u'0x%X (%s)' % \
#                resource_type_name = 'Id'
                resource_type_name = pefile.RESOURCE_TYPE.get(resource_type.struct.Id, 
                                                resource_type.struct.Id)
            resource_type_node = __add_node_attrs_from_structure(resource_type.struct, 
                                      u'Entry', doc, resource_directory_node)
            __set_attr(u'Name', resource_type_name, resource_type_node)
            
            if hasattr(resource_type, u'directory'):
                # PE/IMAGE_DIRECTORY_ENTRY_RESOURCE/Entry/Directory
                resource_type_directory_node = \
                            __add_node_attrs_from_structure(resource_type.directory.struct, 
                                                  u'Directory', 
                                                  doc, resource_type_node)
                for resource_id in resource_type.directory.entries:
                    if resource_id.name is not None:
                        resource_id_name = resource_id.name
                    else:
                        resource_id_name = u'0x%X' % resource_id.struct.Id
                    # PE/IMAGE_DIRECTORY_ENTRY_RESOURCE/Entry/Directory/Entry
                    resource_id_node = __add_node_attrs_from_structure(
                                                            resource_id.struct, 'Entry', 
                                                            doc, resource_type_directory_node)
                    __set_attr(u'Name', resource_id_name, resource_id_node)
                    if hasattr(resource_id, u'directory'):
                        # PE/IMAGE_DIRECTORY_ENTRY_RESOURCE/Entry/Directory/Entry/Directory
                        resource_id_directory_node = __add_node_attrs_from_structure(
                                                        resource_id.directory.struct, 
                                                        u'Directory', 
                                                        doc, resource_id_node)
                        for resource_lang in resource_id.directory.entries:
                            # PE/IMAGE_DIRECTORY_ENTRY_RESOURCE/Entry/Directory/Entry/Directory/Language
                            resource_lang_node = __add_node_attrs_from_structure(
                                                         resource_lang.struct, 
                                                       u'Language', 
                                                       doc, resource_id_directory_node)
                            # PE/IMAGE_DIRECTORY_ENTRY_RESOURCE/Entry/Directory/Entry/Directory/Language/Data
                            if hasattr(resource_lang, u'data'):
                                __add_node_attrs_from_structure(resource_lang.data.struct, 
                                                               u'Data', 
                                                               doc, resource_lang_node)
    # PE/TLS
    if (hasattr(pe, u'DIRECTORY_ENTRY_TLS') and 
         pe.DIRECTORY_ENTRY_TLS and 
         pe.DIRECTORY_ENTRY_TLS.struct):
        tls_node = doc.createElement(u'TLS')
        PE.appendChild(tls_node)
        __add_node_from_structure(pe.DIRECTORY_ENTRY_TLS.struct, 
                                  pe.DIRECTORY_ENTRY_TLS.struct.name, doc, tls_node)

    # PE/Debug
    if hasattr(pe, u'DIRECTORY_ENTRY_DEBUG'):
        debug_node = doc.createElement(u'Debug')
        PE.appendChild(debug_node)
        for dbg in pe.DIRECTORY_ENTRY_DEBUG:
            dbg_node = __add_node_from_structure(dbg.struct, 
                                  dbg.struct.name, doc, debug_node)
            try:
                __add_textnode(pefile.DEBUG_TYPE[dbg.struct.Type], u'Type', doc, dbg_node)
            except KeyError:
                __add_textnode(u'0x%x(Unknown)' % dbg.struct.Type, u'Type', doc, dbg_node)
                
#    # PE/BaseRelocations
#    if hasattr(pe, 'DIRECTORY_ENTRY_BASERELOC'):
#        base_base_relocations_node = doc.createElement('BaseRelocations')
#        PE.appendChild(base_base_relocations_node)
#        
#        for base_reloc in pe.DIRECTORY_ENTRY_BASERELOC:
#            base_reloc_node = __add_node_attrs_from_structure(base_reloc.struct, 
#                                                    None, doc, base_base_relocations_node)
#            base_reloc_node.setAttribute('NumberOfRelocations', str(len(base_reloc.entries)))
#            for reloc in base_reloc.entries:
##                reloc_attrs = {'RVA': '0x%08X' % reloc.rva}
#                reloc_attrs = {'RVA': '%s' % reloc.rva}
#                reloc_type = ''
#                try:
#                    reloc_attrs['Type'] = '%s' % pefile.RELOCATION_TYPE[reloc.type][16:]
#                except KeyError:
#                    reloc_attrs['Type'] = 'Unknown (%d)' % reloc.type
#                __add_node_attrs_from_dict(reloc_attrs, 'Relocation', doc, base_reloc_node)

    # PE/SignedData
    if hasattr(pe, u'SignedData'):
        signedDataNode = pe.SignedData.create_xml_node(pe.SignedData, doc, PE)
# --------------------------xml done. ------------------------------
    try:
        f = codecs.open(xmlpath, u'wb', u'utf-8') 
        try:
            doc.writexml(f, addindent=u'  ', newl=os.linesep, encoding=u'utf-8')
        finally:
            f.close()
    except IOError, e:
        pass
    return pe

from datetime import datetime
def test():
    path = u'e:/softs/pes/'
    path = u'D:/testsamples/todo'
    for filename in os.listdir(path):
        filename = os.path.join(path, filename)
        if os.path.isdir(filename) or os.path.splitext(filename)[1] == u'.txt' \
            or os.path.splitext(filename)[1] == u'.xml':
            continue
        print u'*' * 60
        print filename
        start = datetime.now()
        pe = pe2xml(filename)
        print pe.warnings
        print pe.versioninfo
#        print get_pe_fileinfo(filename, pe)
        break
        
if __name__ == u'__main__':
#    test()
#    print pe2xml(ur'C:\Program Files\MySQL\MySQL Server 5.0\bin\comp-err.exe')
    
    print is_pe(ur'D:\mercury\sepython86\djangoproject\seblog_project\seblog\web\frontapp\media\images\icons.gif')
    print is_pe(ur'D:\mercury\sepython86\djangoproject\seblog_project\seblog\web\frontapp\media\images\base\icons\application.png')
    print get_pe_fileinfo(ur'D:\Program Files\QQfeiche\QQSpeedLauncher.exe')
    